package com.wugh.fast.kernel.util;

import cn.hutool.core.collection.CollectionUtil;
import com.google.common.collect.Lists;
import com.wugh.fast.kernel.util.exception.FastNotRootException;

import java.util.*;
import java.util.function.BiPredicate;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * 说明: 树工具
 *
 * @author wugh
 * @date 2019-04-12 16:18
 */
public class TreeUtils {
    private TreeUtils() {
        //工具类,禁止创建
    }

    /**
     * 通过集合和根节点返回一颗树
     *
     * @param list   node节点集合
     * @param rootId 根节点Id
     * @param <N>    {@linkplain TreeNodeInterface node节点实现}
     * @return 返回Optional, 方便处理空指针
     * @throws FastNotRootException 如果通过rootId得到多个根节点,抛出该异常
     */
    public static <N extends TreeNodeInterface> Optional<N> toTree(List<N> list, Object rootId) throws FastNotRootException {
        List<N> treeList = toList(list, rootId);
        if (treeList.isEmpty()) {
            return Optional.empty();
        }
        if (treeList.size() != 1) {
            //多个结果
            throw new FastNotRootException(rootId, treeList.size());
        }
        return treeList.stream().findFirst();
    }

    /**
     * 通过集合和根节点返回一颗树
     *
     * @param list node节点集合
     * @param <N>  {@linkplain TreeNodeInterface node节点实现}
     * @return 返回Optional, 方便处理空指针
     * @throws FastNotRootException 如果通过rootId得到多个根节点,抛出该异常
     */
    public static <N extends TreeNodeInterface> Optional<N> toTree(List<N> list) throws FastNotRootException {
        List<N> treeList = toList(list);
        if (treeList.isEmpty()) {
            return Optional.empty();
        }
        if (treeList.size() != 1) {
            //多个结果
            throw new FastNotRootException("不指定根节点", treeList.size());
        }
        return treeList.stream().findFirst();
    }

    /**
     * 获取集合和父节点Id
     *
     * @param list     node节点集合
     * @param parentId 父节点Id
     * @param <N>      {@linkplain TreeNodeInterface node节点实现}
     * @return 返回List结合, 没有则是空的集合
     */
    public static <N extends TreeNodeInterface> List<N> toList(List<N> list, Object parentId) {
        return toList(list, Lists.newArrayList(parentId));
    }

    /**
     * 通过已有的集合,自动组装成树,并返回最根部的节点结合
     *
     * @param list node节点集合
     * @param <N>  {@linkplain TreeNodeInterface node节点实现}
     * @return 返回List结合, 没有则是空的集合
     */
    public static <N extends TreeNodeInterface> List<N> toList(List<N> list) {
        return toList(list, Collections.emptyList());
    }

    /**
     * 通过已有的集合,过滤出指定的根节点,并拼接成树
     *
     * @param list node节点集合
     * @param <N>  {@linkplain TreeNodeInterface node节点实现}
     * @return 返回List结合, 没有则是空的集合
     */
    public static <N extends TreeNodeInterface> List<N> toList(List<N> list, Collection<Object> parentIds) {
        Set<Object> ids = list.stream().map(TreeNodeInterface::getId).collect(Collectors.toSet());
        Map<Object, List<N>> pidMap = list.stream().collect(Collectors.groupingBy(TreeNodeInterface::getParentId, LinkedHashMap::new, Collectors.toList()));
        //需要处理的父节点
        List<Object> pids;
        if (CollectionUtil.isEmpty(parentIds)) {
            //找出最根部的pid,即pid,不存在于ids中的
            pids = pidMap.keySet().stream().filter(pid -> !ids.contains(pid)).collect(Collectors.toList());
        } else {
            pids = Lists.newArrayList(parentIds);
        }

        return pids.stream()
                .flatMap(pid -> pidMap.get(pid).stream()
                        .peek(n -> {
                            //补充children
                            loadChildren(n, node -> {
                                if (!pidMap.containsKey(node.getId())) {
                                    //没有子,返回空集合
                                    return Collections.emptyList();
                                }
                                //有子
                                return pidMap.get(node.getId());
                            }, true);
                        }))
                .collect(Collectors.toList());
    }

    /**
     * 加载节点的子节点
     *
     * @param node      node节点
     * @param query     查询方法
     * @param recursive 是否递归子孙节点
     * @param <N>       {@linkplain TreeNodeInterface node节点实现}
     */
    public static <N extends TreeNodeInterface> void loadChildren(N node, Function<N, List<N>> query, boolean recursive) {
        loadChildren(Lists.newArrayList(node), query, recursive);
    }

    /**
     * 加载节点的子节点(递归加载子孙节点)
     *
     * @param node  node节点
     * @param query 查询方法
     * @param <N>   {@linkplain TreeNodeInterface node节点实现}
     */
    public static <N extends TreeNodeInterface> void loadChildren(N node, Function<N, List<N>> query) {
        loadChildren(node, query, true);
    }

    /**
     * 加载节点的子节点
     *
     * @param nodes     node集合
     * @param query     查询方法
     * @param recursive 是否递归子孙节点
     * @param <N>       {@linkplain TreeNodeInterface node节点实现}
     */
    public static <N extends TreeNodeInterface> void loadChildren(List<N> nodes, Function<N, List<N>> query, boolean recursive) {
        Function<Collection<N>, Collection<N>> func = list -> list.stream().flatMap(node -> {
            List<N> children = query.apply(node);
            if(CollectionUtil.isNotEmpty(children)){
                node.setChildren(children);
                return children.stream();
            }else {
                node.setChildren(null);
                return Stream.empty();
            }
        }).collect(Collectors.toList());
        RecursiveUtils.buildSingle(func, (p, n) -> recursive && CollectionUtil.isNotEmpty(n))
                .exec(nodes);
    }

    /**
     * 加载节点的子节点(递归加载子孙节点)
     *
     * @param nodes node集合
     * @param query 查询方法
     * @param <N>   {@linkplain TreeNodeInterface node节点实现}
     */
    public static <N extends TreeNodeInterface> void loadChildren(List<N> nodes, Function<N, List<N>> query) {
        loadChildren(nodes, query, true);
    }

    /**
     * 加载父级数据
     *
     * @param node          节点集合
     * @param query         查询方法
     * @param containMySelf 返回的数据是否包括自身
     * @param recursive     是否递归加载
     * @param <N>           {@linkplain TreeNodeInterface node节点实现}
     * @return 返回查找的结果
     */
    public static <N extends TreeNodeInterface> List<N> addParents(N node, Function<N, N> query, boolean containMySelf, boolean recursive) {
        if (node == null) {
            return new ArrayList<>(0);
        }
        Function<Collection<N>, Collection<N>> listQuery = nodes -> nodes.stream().map(query).filter(Objects::nonNull).collect(Collectors.toList());
        return addParents(Lists.newArrayList(node), listQuery, containMySelf, recursive);
    }

    /**
     * 加载父级数据(包括自身和递归查询父级)
     *
     * @param node  节点集合
     * @param query 单个父节点查询方法
     * @param <N>   {@linkplain TreeNodeInterface node节点实现}
     * @return 返回查找的结果
     */
    public static <N extends TreeNodeInterface> List<N> addParents(N node, Function<N, N> query) {
        return addParents(node, query, true, true);
    }

    /**
     * 加载父级数据
     *
     * @param nodes         节点集合
     * @param query         多个父节点查询方法
     * @param containMySelf 返回的数据是否包括自身
     * @param recursive     是否递归加载
     * @param <N>           {@linkplain TreeNodeInterface node节点实现}
     * @return 返回查找的结果
     */
    public static <N extends TreeNodeInterface> List<N> addParents(Collection<N> nodes, Function<Collection<N>, Collection<N>> query, boolean containMySelf, boolean recursive) {
        if (CollectionUtil.isEmpty(nodes)) {
            return Collections.emptyList();
        }
        if (query == null) {
            return new ArrayList<>(nodes);
        }
        LinkedHashMap<Object, N> idMap4Total = containMySelf ? nodes.stream().collect(Collectors.toMap(TreeNodeInterface::getId, n -> n, (n1, n2) -> n1, LinkedHashMap::new)) : new LinkedHashMap<>();
        BiPredicate<Collection<N>, Collection<N>> predicate = (children, parents) -> {
            if (!recursive || CollectionUtil.isEmpty(parents)) {
                //不递归或者没有父节点,则不进入下次递归
                return false;
            }
            return parents.stream().map(TreeNodeInterface::getParentId).filter(Objects::nonNull).anyMatch(pid -> !idMap4Total.containsKey(pid));
        };
        BinaryOperator<Collection<N>> mapFunc = (children, parents) -> parents.stream().filter(n -> !idMap4Total.containsKey(n.getParentId())).collect(Collectors.toList());
        BinaryOperator<Collection<N>> reduce = (prev, current) -> {
            current.stream().filter(n -> !idMap4Total.containsKey(n.getId())).forEach(n -> idMap4Total.put(n.getId(), n));
            return idMap4Total.values();
        };
        Collection<N> result = RecursiveUtils.buildSingle(query, predicate).map(mapFunc).reduce(Collections.emptyList(), reduce).exec(nodes);
        return Lists.newArrayList(result);
    }

    /**
     * 加载父级数据(包括自身和递归查询父级)
     *
     * @param nodes 节点集合
     * @param query 多个父节点查询方法
     * @param <N>   {@linkplain TreeNodeInterface node节点实现}
     * @return 返回查找的结果
     */
    public static <N extends TreeNodeInterface> List<N> addParents(Collection<N> nodes, Function<Collection<N>, Collection<N>> query) {
        return addParents(nodes, query, true, true);
    }

    /**
     * 通过id去过滤树
     * <pre>
     *     自动保留选中节点的父节点
     * </pre>
     *
     * @param nodes
     * @param ids
     * @param <N>
     * @return
     */
    public static <N extends TreeNodeInterface> List<N> filter(Collection<N> nodes, Collection<?> ids) {
        return filter(nodes, (n) -> ids.contains(n.getId()));
    }

    /**
     * 通过具体方法去过滤树节点
     *
     * @param nodes           节点
     * @param filterPredicate 过滤方法
     * @param <N>
     * @return
     */
    public static <N extends TreeNodeInterface> List<N> filter(Collection<N> nodes, Predicate<N> filterPredicate) {
        List<N> total = RecursiveUtils.exec2Stream(nodes, TreeNodeInterface::getChildren).collect(Collectors.toList());
        Map<Object, N> idMap = total.stream().collect(Collectors.toMap(TreeNodeInterface::getId, n -> n));
        List<N> targets = total.stream().filter(filterPredicate).collect(Collectors.toList());
        List<N> result = addParents(targets, list -> list.stream().map(TreeNodeInterface::getParentId).map(idMap::get).filter(Objects::nonNull).collect(Collectors.toList()));
        result.sort(Comparator.comparing(total::indexOf));
        return toList(result);
    }
}
